#!/usr/bin/env python2
# -*- coding: utf-8 -*-

from __future__ import print_function
from triton import TritonContext, ARCH, MODE, AST_REPRESENTATION, Instruction, OPERAND
from triton import MemoryAccess,CPUSIZE
import sys
import os
import lief

# 加载 elf 文件到内存


INPUT_ADDR = 0x100000

RBP_ADDR = 0x600000
RSP_ADDR = RBP_ADDR - 0x200000


def loadBinary(ctx, path):
    binary = lief.parse(path)
    phdrs = binary.segments
    for phdr in phdrs:
        size = phdr.physical_size
        vaddr = phdr.virtual_address
        print('[+] Loading 0x%06x - 0x%06x' % (vaddr, vaddr+size))
        ctx.setConcreteMemoryAreaValue(vaddr, phdr.content)
    return

if __name__ == '__main__':
    ctx = TritonContext()
    ctx.setArchitecture(ARCH.X86_64)
    ctx.enableMode(MODE.ALIGNED_MEMORY, True)

    loadBinary(ctx, os.path.join(os.path.dirname(__file__), 'crackme_xor'))
    ctx.setAstRepresentationMode(AST_REPRESENTATION.PYTHON)

    pc = 0x0400556

    # 参数是输入字符串的指针
    ctx.setConcreteRegisterValue(ctx.registers.rdi, INPUT_ADDR)

    # 设置栈的值
    ctx.setConcreteRegisterValue(ctx.registers.rsp, RSP_ADDR)
    ctx.setConcreteRegisterValue(ctx.registers.rbp, RBP_ADDR)


    for index in range(5):
            ctx.setConcreteMemoryValue(MemoryAccess(INPUT_ADDR + index, CPUSIZE.BYTE), ord('b'))
            ctx.convertMemoryToSymbolicVariable(MemoryAccess(INPUT_ADDR + index, CPUSIZE.BYTE))

    ast = ctx.getAstContext()
    while pc:
        # Build an instruction
        inst = Instruction()
        opcode = ctx.getConcreteMemoryAreaValue(pc, 16)
        inst.setOpcode(opcode)
        inst.setAddress(pc)

        # 执行指令
        ctx.processing(inst)

        if inst.getAddress() == 0x400597:
            zf   = ctx.getRegisterAst(ctx.registers.zf)
            cstr  = ast.land([
                        ctx.getPathConstraintsAst(),
                        zf == 1
                    ])
            
            # 为暂时求出的解具体化
            model = ctx.getModel(cstr)
            for k, v in list(model.items()):
                value = v.getValue()
                ctx.setConcreteVariableValue(ctx.getSymbolicVariableFromId(k), value)



        if inst.getAddress() == 0x4005B1:
            model = ctx.getModel(ctx.getPathConstraintsAst())
            answer = ""
            for k, v in list(model.items()):
                value = v.getValue()
                answer += chr(value)
            print("answer: {}".format(answer))
            break        


        # 取出下一条指令的地址
        pc = ctx.getConcreteRegisterValue(ctx.registers.rip)

    sys.exit(0)
